import os
import numpy as np
import random
import activations # 导入激活函数


class NeuralNetwork(object):
    def __init__(
            self,
            sizes=[784, 30, 10], # 这里的sizes是输入层到输出层的神经元数量
            #定义了神经网络各层的神经元数量，默认值为 `[784, 30, 10]`，表示输入层有784个神经元，隐藏层有30个神经元，输出层有10个神经元。
            learning_rate=1e-2,
            mini_batch_size=16,
            activation_fn="relu"  # 定义激活函数为relu
    ):
        self.sizes = sizes  # 输入层到输出层的神经元数量
        self.num_layers = len(sizes)  # 层数
        self.activation_fn = getattr(activations, activation_fn) # 获取激活函数
        self.activation_fn_prime = getattr(activations, f"{activation_fn}_prime") # 获取激活函数的导数
        self.weights = [np.array([0])] + [np.random.randn(y, x) / np.sqrt(x) for y, x in
                                          zip(sizes[1:], sizes[:-1])]
        #对于每一层，权重矩阵的大小为 (当前层神经元数量, 前一层神经元数量)，并使用 Xavier 初始化方法（除以 np.sqrt(x)）。
        self.biases = [np.array([0])] + [np.random.randn(y, 1) for y in sizes[1:]]
        #初始化偏置向量。对于每一层，偏置向量的大小为 (当前层神经元数量, 1)，并使用随机初始化。
        self._zs = [np.zeros(bias.shape) for bias in self.biases]   #初始化存储每一层加权输入的中间变量。
        self._activations = [np.zeros(bias.shape) for bias in self.biases] #初始化存储每一层激活值的中间变量。
        self.mini_batch_size = mini_batch_size
        self.lr = learning_rate

    def fit(self, training_data, validation_data=None, epochs=10):
        for epoch in range(epochs):  # 这是训练循环
            random.shuffle(training_data)  # 随机打乱训练数据
            mini_batches = [  # 这是分割数据集的函数
                training_data[k:k + self.mini_batch_size] for k in #将训练数据分割成多个小批量（mini-batches），每个小批量的尺寸由 self.mini_batch_size 决定。
                range(0, len(training_data), self.mini_batch_size)]
            for mini_batch in mini_batches:
                #初始化用于存储偏置和权重的梯度的变量 nabla_b 和 nabla_w。
                nabla_b = [np.zeros(bias.shape) for bias in self.biases]
                print("nabla_b",np.array(nabla_b,dtype=object).shape)
                nabla_w = [np.zeros(weight.shape) for weight in self.weights]
                print("nabla_w",np.array(nabla_w,dtype=object).shape)
                for x, y in mini_batch:
                    #对于每个小批量中的每个样本 (x, y)，进行前向传播和反向传播，计算梯度，并累加到 nabla_b 和 nabla_w 中。
                    self._forward_prop(x)
                    delta_nabla_b, delta_nabla_w = self._back_prop(x, y)
                    print("delta_nabla_b",np.array(delta_nabla_b,dtype=object).shape)
                    #将多个可迭代对象（如列表、元组等）打包成一个元组序列。在这段代码中，zip() 函数被用来同时遍历两个列表，并将它们的对应元素组合在一起
                    nabla_b = [nb + dnb for nb, dnb in zip(nabla_b, delta_nabla_b)] # 累加偏置梯度。
                    nabla_w = [nw + dnw for nw, dnw in zip(nabla_w, delta_nabla_w)] # 累加权重梯度。
                    print("nabla_b_train",np.array(nabla_b,dtype     =object).shape)
                    print("nabla_w_train",np.array(nabla_w ,dtype=object).shape)
                self.weights = [
                    #使用梯度下降法更新权重和偏置。
                    w - (self.lr / self.mini_batch_size) * dw for w, dw in
                    zip(self.weights, nabla_w)]
                self.biases = [
                    b - (self.lr / self.mini_batch_size) * db for b, db in
                    zip(self.biases, nabla_b)]
                print("weights",np.array(self.weights,dtype=object).shape)
                print("biases",np.array(self.biases, dtype=object).shape)
            if validation_data:
                accuracy = self.validate(validation_data) / 100.0  # 计算验证集上的准确率。
                print(f"Epoch {epoch + 1}, accuracy {accuracy} %.")  # 打印训练进度。
            else:
                print(f"Processed epoch {epoch}.")

    def validate(self, validation_data):
        validation_results = [(self.predict(x) == y) for x, y in validation_data]
        #每个元素是一个元组 (x, y)，其中 x 是输入数据，y 是相应的标签。对于每个 (x, y) 对，调用 self.predict(x) 方法进行预测，并将预测结果与真实标签 y 进行比较
        return sum(result for result in validation_results)
    #所有布尔值结果求和。由于在 Python 中，True 被视为 1，False 被视为 0，因此这个求和结果实际上是计算预测正确的样本数量。

    def predict(self, x):
        self._forward_prop(x)  # 前向传播 计算每一层的激活值,
        return np.argmax(self._activations[-1])  # 返回预测结果  获取最后一层激活值中的最大值的索引。这个索引对应于预测的类别。

    def _forward_prop(self, x):
        self._activations[0] = x
        for i in range(1, self.num_layers):
            self._zs[i] = (
                    self.weights[i].dot(self._activations[i - 1]) + self.biases[i]
            )#计算第 i 层的加权输入 self._zs[i]，即前一层激活值 self._activations[i - 1]
            # 与当前层权重 self.weights[i] 的点积，再加上当前层的偏置 self.biases[i]。
            if i == self.num_layers - 1:
                '''如果是最后一层（输出层），使用 softmax 函数作为激活函数。
                    否则，使用在初始化时定义的激活函数 self.activation_fn。'''
                self._activations[i] = activations.softmax(self._zs[i])
            else:

                self._activations[i] = self.activation_fn(self._zs[i])

    def _back_prop(self, x, y):
        nabla_b = [np.zeros(bias.shape) for bias in self.biases]
        nabla_w = [np.zeros(weight.shape) for weight in self.weights]
        error = (self._activations[-1] - y)
        #计算输出层的误差 error，即最后一层的激活值 self._activations[-1] 与真实标签 y 的差。
        nabla_b[-1] = error # 输出层的偏置梯度 nabla_b[-1] 等于误差 error。
        nabla_w[-1] = error.dot(self._activations[-2].transpose()) # 输出层的权重梯度 nabla_w[-1] 等于误差 error 与前一层激活值 self._activations[-2] 的转置的点积。
        for l in range(self.num_layers - 2, 0, -1):
            '''从倒数第二层开始，向前遍历到第一层（隐藏层）。
计算当前层的误差 error，即后一层权重 self.weights[l + 1] 的转置与当前误差的点积，
再乘以当前层激活函数的导数 self.activation_fn_prime(self._zs[l])。
更新当前层的偏置和权重的梯度 nabla_b[l] 和 nabla_w[l]。'''
            error = np.multiply(
                self.weights[l + 1].transpose().dot(error),
                self.activation_fn_prime(self._zs[l])
            )
            nabla_b[l] = error
            nabla_w[l] = error.dot(self._activations[l - 1].transpose())
        return nabla_b, nabla_w     # 返回权重和偏置的梯度。

    def load(self, filename='model.npz'):
        npz_members = np.load(os.path.join(os.curdir, 'models', filename),allow_pickle=True) #允许加载pickle文件作为对象加载到张量
        print("Keys in the npz file:", npz_members.files)  # 打印npz文件中的键值
        self.weights = list(npz_members['weights'])  # 加载权重
        self.biases = list(npz_members['biases'])  # 加载偏置
        self.sizes = [b.shape[0] for b in self.biases]  # 重新设置神经网络的大小
        self.num_layers = len(self.sizes)  # 重新设置层数
        self._zs = [np.zeros(bias.shape) for bias in self.biases]  # 重新初始化中间值
        self._activations = [np.zeros(bias.shape) for bias in self.biases]   # 重新初始化激活值
        self.mini_batch_size = int(npz_members['mini_batch_size'])  # 加载mini_batch_size
        self.lr = float(npz_members['lr'])  # 加载学习率

    def save(self, filename='model.npz'):
        np.savez_compressed(
            file=os.path.join(os.curdir, 'models', filename), # 保存模型的文件夹src/models/model.npz
            weights=np.asarray(self.weights,dtype=object)  ,  # 保存权重
            biases=np.asarray(self.biases,dtype=object),
            mini_batch_size=self.mini_batch_size,  # 保存mini_batch_size
            lr=self.lr
        )
